package com.example.security.config;

import com.example.security.annotation.RateLimit;
import com.example.security.exception.BusinessException;
import lombok.RequiredArgsConstructor;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.List;

/**
 * 接口限流配置
 * 使用Redis+Lua脚本实现分布式限流
 */
@Aspect
@Component
@RequiredArgsConstructor
public class RateLimitConfig {

    private final RedisTemplate<String, Object> redisTemplate;
    
    /**
     * Redis限流Lua脚本
     * 使用令牌桶算法实现限流
     */
    private static final String RATE_LIMIT_SCRIPT = """
        local key = KEYS[1]
        local limit = tonumber(ARGV[1])
        local current = tonumber(redis.call('get', key) or "0")
        if current + 1 > limit then
            return 0
        else
            redis.call("INCRBY", key, 1)
            redis.call("EXPIRE", key, ARGV[2])
            return current + 1
        end
    """;

    /**
     * 限流处理切面
     * 在方法执行前进行限流检查
     */
    @Around("@annotation(rateLimit)")
    public Object rateLimit(ProceedingJoinPoint point, RateLimit rateLimit) throws Throwable {
        // 获取限流key
        String key = "rate_limit:" + point.getSignature().toString();
        
        // 执行限流脚本
        RedisScript<Long> script = new DefaultRedisScript<>(RATE_LIMIT_SCRIPT, Long.class);
        List<String> keys = Collections.singletonList(key);
        
        Long current = redisTemplate.execute(
            script,
            keys,
            rateLimit.limit(),
            rateLimit.timeout()
        );
        
        // 超过限制则抛出异常
        if (current != null && current == 0) {
            throw new BusinessException("请求过于频繁，请稍后重试");
        }
        
        return point.proceed();
    }
} 